import torch
from ultralytics import YOLO
from functools import lru_cache

class YOLODetector:
    _instances = {}

    def __new__(cls):
        gpu_id = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
        if gpu_id not in cls._instances:
            cls._instances[gpu_id] = super(YOLODetector, cls).__new__(cls)
            cls._instances[gpu_id]._initialize(gpu_id)
        return cls._instances[gpu_id]

    def _initialize(self, gpu_id):
        """延迟初始化，避免导入时就加载模型"""
        self._model = None
        self._model_loaded = False
        self._gpu_id = gpu_id

    def load_model(self, model_path="yolo11n.pt"):
        """加载模型（线程安全）"""
        if not self._model_loaded:
            self._model = YOLO(model_path).float()
            self._model.eval()
            if self._gpu_id != 'cpu':
                self._model = self._model.to(f"cuda:{self._gpu_id}")
            self._model_loaded = True

    @torch.no_grad()
    def detect_boxes(self, images, conf=0.5):
        """
        检测图像中的边界框
        Args:
            images: Tensor [B, C, H, W] 或图像路径列表
            conf: 置信度阈值
        Returns:
            List[torch.Tensor]: 每个图像的边界框 [N, 4] (x1, y1, x2, y2)
        """
        if not self._model_loaded:
            self.load_model()

        results = self._model(images, conf=conf, verbose=False)

        boxes_list = []
        for result in results:
            if result.boxes is not None and len(result.boxes) > 0:
                boxes = result.boxes.xyxy  # [N, 4]
                boxes_list.append(boxes.cpu())
            else:
                # 没有检测到目标，返回空tensor
                boxes_list.append(torch.empty(0, 4))

        return boxes_list

# 全局实例
yolo_detector = YOLODetector()